{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import orchest\n",
    "import orchest.parameters as op\n",
    "import redis\n",
    "import psycopg2\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "step_params = op.get_params()[0]\n",
    "internal_connectivity = dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test connectivity of all redis services.\n",
    "redis_services = step_params.get(\"redis-services\", [\"redis\"])\n",
    "internal_connectivity[\"redis-services\"] = defaultdict(bool)\n",
    "for svc in redis_services:\n",
    "    try:\n",
    "        redis_host = orchest.get_service(svc)[\"internal_hostname\"]\n",
    "        redis_client = redis.Redis(host=redis_host, port=6379, db=0)\n",
    "        redis_client.set(\"hello\", \"there\")\n",
    "        if redis_client.get(\"hello\") == b\"there\":\n",
    "            internal_connectivity[\"redis-services\"][svc] = True\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test connectivity of all postgres services.\n",
    "postgres_services = step_params.get(\"postgres-services\", [\"postgres\"])\n",
    "internal_connectivity[\"postgres-services\"] = defaultdict(bool)\n",
    "for svc in postgres_services:\n",
    "    try:\n",
    "        postgres_host = orchest.get_service(\"postgres\")[\"internal_hostname\"]\n",
    "        conn = psycopg2.connect(dbname=\"postgres\", user=\"postgres\", host=postgres_host)\n",
    "        cur = conn.cursor()\n",
    "        cur.execute(\"CREATE TABLE test (id serial PRIMARY KEY, num integer, data varchar);\")\n",
    "        cur.execute(\"INSERT INTO test (num, data) VALUES (%s, %s)\", (1337, \"hello\"))\n",
    "        cur.execute(\"SELECT * FROM test;\")\n",
    "        if cur.fetchone() == (1, 1337, \"hello\"):\n",
    "            internal_connectivity[\"postgres-services\"][svc] = True\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "external_connectivity = defaultdict(list)\n",
    "for name, service in orchest.get_services().items():\n",
    "    for url in service[\"external_urls\"].values():\n",
    "        external_connectivity[name].append(url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_file_path = step_params.get(\"test-output-file\", \"test-output.json\")\n",
    "connectivity = {\n",
    "    \"internal-connectivity\": internal_connectivity,\n",
    "    \"external-connectivity\": external_connectivity,\n",
    "}\n",
    "# Internal connectivity maps service type (redis/postgres) to a dict\n",
    "# mapping service name to the result of the connectivity check (True,\n",
    "# False). External connectivity maps service name to a list of\n",
    "# external urls where host and port need to be replaced by the host and\n",
    "# port on which Orchest is running.\n",
    "with open(f\"/data/{output_file_path}\", \"w\") as output_file:\n",
    "    json.dump(connectivity, output_file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "orchest-kernel-30cfd18d-c5ec-4651-8a9b-6e12c36fc3b5"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
